import os
import random

import cv2
import json
from tqdm import tqdm
import shutil


def load_spire_dir(spire_dir):
    """
    :param spire_dir: path to spire dataset root dir
    :return: List - img_list, List - ann_list
    """
    img_list = []
    ann_list = []
    for sub_dir in os.listdir(spire_dir):
        sub_fn = os.path.join(spire_dir, sub_dir)
        if os.path.isdir(sub_fn):
            if sub_dir == 'scaled_images':
                img_list.extend([os.path.join(sub_fn, img_nm) for img_nm in os.listdir(sub_fn) if
                                 img_nm.endswith('.jpg') or img_nm.endswith('.png')])
                ann_list.extend([os.path.join(spire_dir, 'annotations', img_nm + '.json') for img_nm in
                                 os.listdir(sub_fn) if
                                 img_nm.endswith('.jpg') or img_nm.endswith('.png')])
            else:
                for sub_sub_dir in os.listdir(sub_fn):
                    sub_sub_fn = os.path.join(sub_fn, sub_sub_dir)
                    if os.path.isdir(sub_sub_fn):
                        if sub_sub_dir == 'scaled_images':
                            img_list.extend([os.path.join(sub_sub_fn, img_nm) for img_nm in os.listdir(sub_sub_fn) if
                                             img_nm.endswith('.jpg') or img_nm.endswith('.png')])
                            ann_list.extend([os.path.join(sub_fn, 'annotations', img_nm + '.json') for img_nm in
                                             os.listdir(sub_sub_fn) if
                                             img_nm.endswith('.jpg') or img_nm.endswith('.png')])
    return img_list, ann_list


def format_check(spire_root_dir, spire_train_dir, spire_val_dir):
    img_list, ann_list = load_spire_dir(spire_root_dir)
    all_cates = []
    num_cates = {}
    orin_basename = os.path.basename(spire_root_dir)
    train_basename = os.path.basename(spire_train_dir)
    val_basename = os.path.basename(spire_val_dir)

    for i in tqdm(range(len(img_list))):
        img_fn = img_list[i]
        ann_fn = ann_list[i]

        img = cv2.imread(img_fn)
        ann = json.loads(open(ann_fn, 'r').read())

        if random.random() < 0.8:
            img_fn2 = img_fn.replace(orin_basename, train_basename)
            ann_fn2 = ann_fn.replace(orin_basename, train_basename)
            img_fn2_dn = os.path.dirname(img_fn2)
            ann_fn2_dn = os.path.dirname(ann_fn2)
            if not os.path.exists(img_fn2_dn):
                os.makedirs(img_fn2_dn)
            if not os.path.exists(ann_fn2_dn):
                os.makedirs(ann_fn2_dn)
            shutil.copy(img_fn, img_fn2)
            shutil.copy(ann_fn, ann_fn2)
            # print(img_fn2)
        else:
            img_fn2 = img_fn.replace(orin_basename, val_basename)
            ann_fn2 = ann_fn.replace(orin_basename, val_basename)
            img_fn2_dn = os.path.dirname(img_fn2)
            ann_fn2_dn = os.path.dirname(ann_fn2)
            if not os.path.exists(img_fn2_dn):
                os.makedirs(img_fn2_dn)
            if not os.path.exists(ann_fn2_dn):
                os.makedirs(ann_fn2_dn)
            shutil.copy(img_fn, img_fn2)
            shutil.copy(ann_fn, ann_fn2)
            # print(img_fn2)

        for a in ann['annos']:
            if a['category_name'] in num_cates:
                num_cates[a['category_name']] += 1
            else:
                num_cates[a['category_name']] = 1
            if a['category_name'] not in all_cates:
                all_cates.append(a['category_name'])

    print('Check Done! Success!')
    print(all_cates)
    print(num_cates)


if __name__ == '__main__':
    spire_root = r'C:\dataset\BB-Air2AirDet-trainval-v230503'
    spire_train_root = r'C:\dataset\BB-Air2AirDet-train-v230512'
    spire_val_root = r'C:\dataset\BB-Air2AirDet-val-v230512'
    format_check(spire_root, spire_train_root, spire_val_root)
